import os

import  requests

def get_single_result(item):
    text = item["prompt"]
    url='http://localhost:8000/v1/chat/completions'
    headers = {"Content-Type":"application/json"}
    request_json = {
  "model": "string",
  "messages": [
    {
      "role": "user",
      "content": text
    }
  ],
  "temperature": 0,
  "top_p": 0,
  "n": 1,
  "max_tokens": 0,
  "stream": False
}
    output = ""
    count = 0
    while count < 3:
        count += 1
        req = requests.post(url,headers=headers,json=request_json)
        if req.json():
            re = req.json()
            output = re["choices"][0]["message"]["content"]
            break

    return output

import pandas as pd



def check_fold(fold_name):
    for root,fold,files in os.walk(fold_name):
        for file in files:
            df = pd.read_json(file,orient="records",lines=True,encoding='gbk')

            df['mpt-30b_ans'] = df.apply(get_single_result,axis=1)


            df.to_json(file,orient="records",lines=True)




if __name__ == "__main__":
    check_fold(r'/root/qiufengfeng/data')
